import scipy.io as sio
import os
import shutil
from PIL import Image

# 数据集图片所在的基础目录
base_image_dir = 'D:/Python_study/Digital_Image_Processing_2.0/stanford_cars/car_ims'

# 加载.mat标注文件
data = sio.loadmat('./stanford_cars/cars_annos.mat')

# 提取标注信息
annotations = data['annotations'][0]

# 创建符合YOLO格式的目录结构
if not os.path.exists('yolo_dataset'):
    os.makedirs('yolo_dataset')
    os.makedirs('yolo_dataset/images/train')
    os.makedirs('yolo_dataset/images/val')
    os.makedirs('yolo_dataset/labels/train')
    os.makedirs('yolo_dataset/labels/val')

# 遍历每一条标注信息
for ann in annotations:
    img_name = ann[0][0]
    bbox = [ann[1][0][0], ann[2][0][0], ann[3][0][0], ann[4][0][0]]
    class_id = 0
    is_test = int(ann[6][0][0])

    # 根据is_test的值确定图像和标签要存放的目标目录
    img_dest_dir = 'yolo_dataset/images/train' if is_test == 0 else 'yolo_dataset/images/val'
    label_dest_dir = 'yolo_dataset/labels/train' if is_test == 0 else 'yolo_dataset/labels/val'

    # 拼接得到图片在原始数据集里的完整路径
    full_img_path = os.path.join(base_image_dir, os.path.basename(img_name))

    # 检查原始路径中图片是否存在，存在则复制到目标图像目录，不存在则打印警告并跳过
    if os.path.exists(full_img_path):
        shutil.copy(full_img_path, os.path.join(img_dest_dir, os.path.basename(img_name)))
    else:
        print(f"警告：图片未在 {full_img_path} 找到，跳过 {img_name}")
        continue

    # 打开图片，获取图片的宽和高，用于后续将边界框坐标归一化
    img = Image.open(full_img_path)
    img_width, img_height = img.size

    # 将边界框的坐标转换为YOLO格式所需的归一化中心坐标和宽高
    x_center = (bbox[0] + bbox[2]) / 2 / img_width
    y_center = (bbox[1] + bbox[3]) / 2 / img_height
    width = (bbox[2] - bbox[0]) / img_width
    height = (bbox[3] - bbox[1]) / img_height

    label_path = os.path.join(label_dest_dir, f'{os.path.splitext(os.path.basename(img_name))[0]}.txt')
    with open(label_path, 'w') as f:
        f.write(f'{class_id} {x_center} {y_center} {width} {height}\n')


print("转换为 YOLO 格式并划分训练集和验证集完成。")